{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "import os\n",
    "import torch\n",
    "import json\n",
    "from tqdm.auto import tqdm\n",
    "from internal.dataparsers.colmap_dataparser import Colmap, ColmapDataParser\n",
    "from internal.utils.gaussian_model_loader import GaussianModelLoader\n",
    "from internal.utils.light_gaussian import get_count_and_score, calculate_v_imp_score, get_prune_mask"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "torch.autograd.set_grad_enabled(False)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "253d162b56dc1681",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "partition_base_dir = os.path.expanduser(\n",
    "    \"~/data/image_set/JNUCar_undistorted/colmap/drone/dense_max_2048/0/partitions-threshold_0.2/\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "58439cf23b16df07",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "dataparser_outputs = ColmapDataParser(\n",
    "    os.path.join(partition_base_dir, \"..\"),\n",
    "    output_path=os.getcwd(),\n",
    "    global_rank=0,\n",
    "    params=Colmap(\n",
    "        appearance_groups=\"appearance_image_dedicated\",\n",
    "        eval_step=32\n",
    "    ),\n",
    ").get_outputs()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2da461e704fd3d02",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "image_name_to_camera_idx = {}\n",
    "for idx, name in enumerate(dataparser_outputs.train_set.image_names):\n",
    "    image_name_to_camera_idx[name] = idx\n",
    "len(image_name_to_camera_idx)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "acd30869a8b7c042",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def get_cameras_by_image_list(image_list: list):\n",
    "    cameras = []\n",
    "    for i in image_list:\n",
    "        cameras.append(dataparser_outputs.train_set.cameras[image_name_to_camera_idx[i]])\n",
    "    return cameras"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "924272b43abdc20",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "partitions = torch.load(os.path.join(partition_base_dir, \"partitions.pt\"),\n",
    "                        map_location=\"cpu\")\n",
    "\n",
    "orientation_transformation = partitions[\"orientation_transformation\"]\n",
    "\n",
    "model_paths = []\n",
    "for idx, i in enumerate(partitions[\"ids\"]):\n",
    "    if len(partitions[\"image_indices\"][idx]) < 32:\n",
    "        continue\n",
    "    model_paths.append((i, os.path.join(\n",
    "        \"../outputs/JNUAerial-0526/\",\n",
    "        f\"P_{i[0]:03d}_{i[1]:03d}.txt\")))\n",
    "partitions.keys(), model_paths"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "aeb88dc2829a6b0c",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "partition_id_to_index = {i: idx for idx, i in enumerate(partitions[\"ids\"])}\n",
    "partition_id_to_index"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e1e8bd2b3bdeee05",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "dtype = torch.float\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "prune_percent = 0.6\n",
    "\n",
    "n_gaussians_before_pruning = 0\n",
    "n_gaussians_after_pruning = 0\n",
    "\n",
    "with tqdm(model_paths) as t:\n",
    "    for i in t:\n",
    "        # if i[0] != (0, 0):\n",
    "        #     continue\n",
    "\n",
    "        partition_id = i[0]\n",
    "        model_output_path = i[1]\n",
    "\n",
    "        partition_xy = partitions[\"xys\"][partition_id_to_index[i[0]]]\n",
    "        load_file = GaussianModelLoader.search_load_file(i[1])\n",
    "        t.set_description(f\"{partition_xy}: {load_file}\")\n",
    "\n",
    "        # load ckpt file\n",
    "        ckpt = torch.load(load_file, map_location=device)\n",
    "        # instantiate GaussianModel\n",
    "        gaussian_model = GaussianModelLoader.initialize_model_from_checkpoint(\n",
    "            ckpt,\n",
    "            device,\n",
    "        )\n",
    "\n",
    "        # get the position of all the Gaussians\n",
    "        xyz = gaussian_model.get_means()\n",
    "\n",
    "        # find those inside bounding box, pruning only apply to them\n",
    "        reoriented_xyz = xyz @ orientation_transformation[:3, :3].T\n",
    "        # include min bound, exclude max bound\n",
    "        is_in_partition = torch.logical_and(torch.ge(reoriented_xyz[:, :2], partition_xy),\n",
    "                                            torch.lt(reoriented_xyz[:, :2], partition_xy + 2 * partitions[\"radius\"]))\n",
    "        is_in_partition = torch.logical_and(is_in_partition[:, 0], is_in_partition[:, 1])\n",
    "\n",
    "        # split into the inside and outside set\n",
    "        outside_gaussians = {}\n",
    "        inside_gaussians = {}\n",
    "        for k, v in gaussian_model.properties.items():\n",
    "            outside_gaussians[k] = v[~is_in_partition]\n",
    "            inside_gaussians[k] = v[is_in_partition]\n",
    "\n",
    "        # update gaussian model, containing only inside set\n",
    "        gaussian_model.properties = inside_gaussians\n",
    "        # move gaussian model to cuda\n",
    "        gaussian_model.cuda()\n",
    "\n",
    "        n_gaussians_before_pruning += gaussian_model.n_gaussians\n",
    "\n",
    "        # get partition image list\n",
    "        with open(os.path.join(model_output_path, \"cameras.json\"), \"r\") as f:\n",
    "            cameras_json = json.load(f)\n",
    "        image_list = [i[\"img_name\"] for i in cameras_json]\n",
    "\n",
    "        cameras = get_cameras_by_image_list(image_list)\n",
    "\n",
    "        # calculate scores\n",
    "        hit_camera_count, opacity_score, alpha_score, visibility_score = get_count_and_score(\n",
    "            gaussian_model,\n",
    "            cameras,\n",
    "            anti_aliased=True,\n",
    "        )\n",
    "\n",
    "        # prune by visibility\n",
    "        # # get prune indices\n",
    "        # visibility_score_close_to_zero = torch.isclose(visibility_score, torch.tensor(0.).to(visibility_score))\n",
    "        # visibility_score_close_to_zero_count = visibility_score_close_to_zero.sum()\n",
    "        # prune_percent = 0.9\n",
    "        # # ignore the Gaussians visibility score close zero\n",
    "        # keep_count = ((visibility_score.shape[0] - visibility_score_close_to_zero_count) * (1 - prune_percent)).to(torch.int)\n",
    "        # # get the indices (in partition) to be preserved\n",
    "        # visibility_score_sorted = torch.sort(visibility_score, descending=True)\n",
    "        # local_indices_to_preserved = visibility_score_sorted.indices[:keep_count].cpu()\n",
    "\n",
    "        # prune by opacity\n",
    "        v_imp_score = calculate_v_imp_score(gaussian_model.get_scaling, opacity_score, 0.1)\n",
    "        prune_mask = get_prune_mask(prune_percent, v_imp_score)\n",
    "        # local_indices_to_preserved = (~prune_mask).nonzero().squeeze(-1).cpu()\n",
    "\n",
    "        # prune gaussian model with the LightGaussian mask\n",
    "        gaussian_model.properties = {k: v[~prune_mask] for k, v in gaussian_model.properties.items()}\n",
    "        n_gaussians_after_pruning += gaussian_model.n_gaussians\n",
    "        # # the indices (in partition) to be pruned\n",
    "        # local_indices_to_prune = visibility_score_sorted.indices[keep_count:]\n",
    "        # is_in_partition_indices = is_in_partition.nonzero().squeeze(-1)\n",
    "        # gaussian_indices_to_prune = is_in_partition_indices[local_indices_to_prune.to(device=is_in_partition_indices.device)]\n",
    "        # \n",
    "        # # convert prune indices to preserve mask\n",
    "        # preserve_mask = torch.ones_like(is_in_partition, dtype=torch.bool)\n",
    "        # preserve_mask[gaussian_indices_to_prune] = False\n",
    "\n",
    "        # update state_dict of checkpoint\n",
    "        \"\"\"\n",
    "        [NOTE]\n",
    "        The codes related to the `static part` below have not been released yet.\n",
    "        So, rather than move some Gaussian to `static part`, you should prune the `state_dict`, `optimizer_states` and `gaussian_model_extra_state_dict` of `ckpt` according to the mask `local_indices_to_preserved`.\n",
    "        \"\"\"\n",
    "        # for i in state_dict:\n",
    "        #     # move those Gaussian outside the partition to static part, which will not be optimized during finetune\n",
    "        #     static_gaussian_property_key = f\"static_{i}\"\n",
    "        #     static_gaussian_property = ckpt[\"state_dict\"][i][~is_in_partition]\n",
    "        #     if static_gaussian_property_key in ckpt[\"state_dict\"]:\n",
    "        #         original_static_gaussian_num = ckpt[\"state_dict\"][static_gaussian_property_key].shape[0]\n",
    "        #         new_static_gaussian_num = static_gaussian_property.shape[0]\n",
    "        #         static_gaussian_property = torch.concat([\n",
    "        #             ckpt[\"state_dict\"][static_gaussian_property_key],\n",
    "        #             static_gaussian_property,\n",
    "        #         ], dim=0)\n",
    "        #         print(f\"#{partition_id}: {original_static_gaussian_num} static Gaussians exists, merge with {new_static_gaussian_num} new static Gaussians, total {static_gaussian_property.shape[0]} after merging\")\n",
    "        #     ckpt[\"state_dict\"][static_gaussian_property_key] = static_gaussian_property\n",
    "        #     # make optimizable Gaussians only contains those locating in partition\n",
    "        #     ckpt[\"state_dict\"][i] = state_dict[i]\n",
    "\n",
    "        # merge inside and outside\n",
    "        gaussian_model.to(device=device)\n",
    "        gaussian_model.properties = {k: torch.concat(\n",
    "            [outside_gaussians[k], v],\n",
    "            dim=0,\n",
    "        ) for k, v in gaussian_model.properties.items()}\n",
    "        \n",
    "        # update gaussian states of ckpt\n",
    "        for k, v in gaussian_model.state_dict().items():\n",
    "            state_dict_full_key = \"gaussian_model.{}\".format(k)\n",
    "            assert state_dict_full_key in ckpt[\"state_dict\"]\n",
    "            ckpt[\"state_dict\"][state_dict_full_key] = v\n",
    "\n",
    "        # prune optimizer state\n",
    "        optimizer_prunable_states = [\"exp_avg\", \"exp_avg_sq\"]\n",
    "        property_names = list(gaussian_model.property_names)\n",
    "        new_optimizer_state_list = []\n",
    "        for optimizer_state in ckpt[\"optimizer_states\"]:\n",
    "            if len(property_names) == 0:\n",
    "                break\n",
    "\n",
    "            for param_group_idx, param_group in enumerate(optimizer_state[\"param_groups\"]):\n",
    "                # whether a gaussian param_group\n",
    "                param_group_name = param_group[\"name\"]\n",
    "                if param_group_name not in property_names:\n",
    "                    continue\n",
    "\n",
    "                outside_gaussian_states = {\n",
    "                    k: optimizer_state[\"state\"][param_group_idx][k][~is_in_partition]\n",
    "                    for k in optimizer_prunable_states\n",
    "                }\n",
    "                inside_gaussian_states = {\n",
    "                    k: optimizer_state[\"state\"][param_group_idx][k][is_in_partition][~prune_mask.to(device=device)]\n",
    "                    for k in optimizer_prunable_states\n",
    "                }\n",
    "\n",
    "                # replace optimizer states\n",
    "                for k in optimizer_prunable_states:\n",
    "                    optimizer_state[\"state\"][param_group_idx][k] = torch.concat(\n",
    "                        [outside_gaussian_states[k], inside_gaussian_states[k]],\n",
    "                        dim=0,\n",
    "                    )\n",
    "\n",
    "                property_names.remove(param_group_name)\n",
    "\n",
    "        # prune density controller state_dict by simply replacing with zeros\n",
    "        for i in ckpt[\"state_dict\"]:\n",
    "            if i.startswith(\"density_controller.\"):\n",
    "                ckpt[\"state_dict\"][i] = torch.zeros((gaussian_model.n_gaussians, *ckpt[\"state_dict\"][i].shape[1:]))\n",
    "\n",
    "        # save checkpoint\n",
    "        checkpoint_save_dir = os.path.join(model_output_path, \"pruned_checkpoints\")\n",
    "        os.makedirs(checkpoint_save_dir, exist_ok=True)\n",
    "        torch.save(ckpt, os.path.join(checkpoint_save_dir, f\"latest-opacity_pruned-{prune_percent}.ckpt\"))\n",
    "        \n",
    "        del ckpt\n",
    "        del gaussian_model\n",
    "        del outside_gaussians\n",
    "        del inside_gaussians\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "f\"{n_gaussians_after_pruning} / {n_gaussians_before_pruning}\""
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e4970f80b12cb03d",
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
